/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.
 * MemCache_Hybrid is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
*/
#ifndef MEM_FABRIC_MMC_LOCALITY_STRATEGY_H
#define MEM_FABRIC_MMC_LOCALITY_STRATEGY_H

#include <vector>
#include <set>
#include <random>
#include <algorithm>
#include <ostream>
#include <cstdint>
#include "mmc_mem_blob.h"
#include "mmc_types.h"
#include "mmc_blob_allocator.h"
#include "mmc_msg_base.h"
#include "mmc_msg_client_meta.h"

namespace ock {
namespace mmc {

enum AllocFlags {
    ALLOC_ARRANGE = 0,
    ALLOC_FORCE_BY_RANK = 1 << 0, // 按照rank强制分配
    ALLOC_RANDOM = 1 << 1,
};

struct MmcLocalMemCurInfo {
    uint64_t capacity_;
};

using MmcMemPoolCurInfo = std::map<MmcLocation, MmcLocalMemCurInfo>;

using MmcAllocators = std::map<MmcLocation, MmcBlobAllocatorPtr>;

class MmcLocalityStrategy : public MmcReferable {
public:
    static Result ArrangeLocality(const MmcAllocators &allocators, const AllocOptions &allocReq,
                                  std::vector<MmcMemBlobPtr> &blobs, std::unordered_set<uint32_t> &excludeRanks)
    {
        if (allocators.empty()) {
            MMC_LOG_ERROR("Cannot allocate blob, allocators empty");
            return MMC_ERROR;
        }
        if (allocators.size() < allocReq.numBlobs_) {
            MMC_LOG_ERROR("Cannot allocate blob, allocators not enough");
            return MMC_ERROR;
        }
        size_t allocatedCount = 0;
        if (!allocReq.preferredRank_.empty() && excludeRanks.find(allocReq.preferredRank_[0]) == excludeRanks.end()) {
            MmcLocation location{};
            location.mediaType_ = static_cast<MediaType>(allocReq.mediaType_);
            location.rank_ = allocReq.preferredRank_[0];
            auto itPrefer = allocators.find(location);
            if (itPrefer != allocators.end()) {
                auto allocator = itPrefer->second;
                MmcMemBlobPtr blob = allocator->Alloc(allocReq.blobSize_);
                if (blob != nullptr) {
                    excludeRanks.insert(location.rank_);
                    blobs.push_back(blob);
                    allocatedCount++;
                    if (allocatedCount == allocReq.numBlobs_) {
                        return MMC_OK;
                    }
                }
            }
        }

        AllocOptions tmpAllocReq = allocReq;
        tmpAllocReq.numBlobs_ = allocReq.numBlobs_ - allocatedCount;
        tmpAllocReq.flags_ = static_cast<uint32_t>(ALLOC_RANDOM);

        return RandomAssign(allocators, tmpAllocReq, blobs, excludeRanks);
    }

    static Result ForceAssign(const MmcAllocators &allocators, const AllocOptions &allocReq,
                                  std::vector<MmcMemBlobPtr> &blobs)
    {
        if (allocators.empty()) {
            MMC_LOG_ERROR("Cannot allocate blob, allocators empty");
            return MMC_ERROR;
        }
        if (allocReq.numBlobs_ != 1) {
            MMC_LOG_ERROR("force assign allocate blob numBlobs must be 1");
            return MMC_ERROR;
        }
        MmcLocation location{};
        location.mediaType_ = static_cast<MediaType>(allocReq.mediaType_);
        location.rank_ = allocReq.preferredRank_.empty() ? 0 : allocReq.preferredRank_[0];
        auto itPrefer = allocators.find(location);
        if (itPrefer == allocators.end()) {
            MMC_LOG_ERROR("Cannot force assign allocate blob, allocator rank: " << location.rank_
                << " media type: " << location.mediaType_ << " not found");
            return MMC_ERROR;
        }
        auto it = itPrefer;

        auto allocator = it->second;
        if (allocator == nullptr) {
            MMC_LOG_ERROR("Cannot force assign allocate blob, allocator rank: " << location.rank_
                << " media type: " << location.mediaType_ << " is nullptr");
            return MMC_ERROR;
        }
        MmcMemBlobPtr blob = allocator->Alloc(allocReq.blobSize_);
        if (blob != nullptr) {
            blobs.push_back(blob);
        } else {
                MMC_LOG_ERROR("Cannot force assign allocate blob, allocator rank: " << location.rank_
                << " media type: " << location.mediaType_ << " cannot alloc space");
            return MMC_ERROR;
        }

        return MMC_OK;
    }

    static Result RandomAssign(const MmcAllocators &allocators, const AllocOptions &allocReq,
                               std::vector<MmcMemBlobPtr> &blobs, std::unordered_set<uint32_t> &excludeRanks)
    {
        if (allocators.empty()) {
            MMC_LOG_ERROR("Cannot allocate blob, allocators empty");
            return MMC_ERROR;
        }
        MmcLocation lowerBound;
        MmcLocation upperBound;
        lowerBound.mediaType_ = static_cast<MediaType>(allocReq.mediaType_);
        lowerBound.rank_ = std::numeric_limits<uint32_t>::min();
        upperBound.mediaType_ = static_cast<MediaType>(allocReq.mediaType_);
        upperBound.rank_ = std::numeric_limits<uint32_t>::max();
        auto first = allocators.lower_bound(lowerBound);
        auto last = allocators.upper_bound(upperBound);
        const size_t numCandidates = static_cast<size_t>(std::distance(first, last));
        if (numCandidates == 0) {
            MMC_LOG_ERROR("No allocators found for media type: " << allocReq.mediaType_);
            return MMC_ERROR;
        }
        if (numCandidates < allocReq.numBlobs_) {
            MMC_LOG_ERROR("Not enough allocators for media type: " << allocReq.mediaType_
                          << ", required: " << allocReq.numBlobs_ << ", available: " << numCandidates);
            return MMC_ERROR;
        }
        std::vector<size_t> indices;
        indices.reserve(numCandidates);
        for (size_t i = 0; i < numCandidates; ++i) {
            indices.push_back(i);
        }
        std::random_device rd;
        std::mt19937 rng(rd());
        std::shuffle(indices.begin(), indices.end(), rng);
        size_t allocatedCount = 0;
        for (size_t i = 0; i < numCandidates; ++i) {
            auto it = first;
            std::advance(it, indices[i]);
            if (excludeRanks.find(it->first.rank_) != excludeRanks.end()) {
                continue;
            }
            MmcMemBlobPtr blob = it->second->Alloc(allocReq.blobSize_);
            if (blob != nullptr) {
                blobs.push_back(blob);
                allocatedCount++;
                if (allocatedCount >= allocReq.numBlobs_) {
                    break;
                }
            }
        }
        if (allocatedCount < allocReq.numBlobs_) {
            MMC_LOG_ERROR("RandomAssign failed to allocate all blobs. Requested: " << allocReq.numBlobs_);
            return MMC_ERROR;
        }
        return MMC_OK;
    }
};

} // namespace mmc
} // namespace ock

#endif // MEM_FABRIC_MMC_LOCALITY_STRATEGY_H